import os

import cv2
from tqdm import tqdm

from option import parse_args


def cmp_video_name(video_name):
    video_name = video_name[:-4]
    video_name_splt = video_name.split('_')
    if len(video_name_splt) == 1:
        return int(video_name[1])
    elif len(video_name_splt) == 2:
        return int(video_name_splt[0][2:]) + 300
    else:
        return int(video_name_splt[0][2:]) * 1000 + int(video_name_splt[1][2:]) * 10 + int(video_name_splt[2])


def unaligned_videos(real, synthesis, func, template: str):
    video_names = os.listdir(synthesis)
    video_names.sort(key=cmp_video_name)
    for video_name in tqdm(video_names):
        video_path_split = video_name.split('_')
        assert len(video_path_split) == 3
        original = '{}_{}'.format(video_path_split[0], video_path_split[2])

        capture = cv2.VideoCapture(os.path.join(synthesis, video_name))
        fake_video_frames_num = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
        capture = cv2.VideoCapture(os.path.join(real, original))
        real_video_frames_num = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
        if fake_video_frames_num != real_video_frames_num:
            func(template.format(video_name))


def main():
    args = parse_args()
    synthesis = os.path.join(args.root_dir, 'Celeb-synthesis')
    real = os.path.join(args.root_dir, 'Celeb-real')
    args.save = True
    if args.save:
        with open('unaligned_videos_statistics.dat', 'w') as f:
            unaligned_videos(real, synthesis, f.write, '{}\n')
    else:
        unaligned_videos(real, synthesis, print, '{}')


if __name__ == '__main__':
    main()
